MODULE pmd1_main

   USE md1_global_consts
   USE md1_global_types
   USE pfmm_global_types
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: pmd1_initial,                        &
             pmd1_final,                          &
             pmd1_write_multipole_integrals,      &
             pmd1_get_J_matrix

   TYPE(md1_basis), SAVE, POINTER :: basis
   REAL(REALK), SAVE :: lattice_constants(3)

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_initial(filename,nbas,lattice_consts)

      USE md1_basis_set, ONLY: md1_init_basis_set, md1_get_basis_set

      IMPLICIT NONE
      CHARACTER(*),  INTENT(IN) :: filename    
      INTEGER(INTK), INTENT(IN) :: nbas
      REAL(REALK),   INTENT(IN) :: lattice_consts(3)

      WRITE(IO1,*)
      WRITE(IO1,*), "------------------------------------------"
      WRITE(IO1,*), "|      Periodic MD1 initialization       |"
      WRITE(IO1,*), "------------------------------------------"
      WRITE(IO1,*)
      NULLIFY(basis)
      CALL md1_init_basis_set(filename)
      CALL md1_get_basis_set(basis)
      IF (nbas /= basis%nbas) STOP 'pmd1_initial: nbas error!'
      lattice_constants(:) = lattice_consts(:)

   END SUBROUTINE pmd1_initial

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_final

      USE md1_basis_set,   ONLY: md1_free_basis_set

      IMPLICIT NONE
      CALL md1_free_basis_set
      NULLIFY(basis)

   END SUBROUTINE pmd1_final

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_get_J_matrix(dens,TQvecs,Jmat)

      USE pmd1_shell_pairs,   ONLY: pmd1_get_shell_pairs

      IMPLICIT NONE
      REAL(REALK),        INTENT(IN)  :: dens(:,:,:)
      TYPE(T_vector_map), INTENT(IN)  :: TQvecs(:)
      REAL(REALK),        INTENT(OUT) :: Jmat(:,:,:)

      INTEGER(INTK), PARAMETER :: iSmax = 6

      TYPE(md1_sh_pairs), POINTER :: mu_nu(:), ls(:)
      INTEGER(INTK) :: t,q, i,j,k
      REAL(REALK) :: mu_vec(3), Tvec(3), box(3)
      REAL(REALK) :: Svec(3), Qvec(3)
      REAL(REALK) :: energy

      IF (SIZE(Jmat,1) > 1) CALL fmm_quit('many AOs NYI')
      IF (SIZE(Jmat,2) > 1) CALL fmm_quit('many AOs NYI')

      energy = zero
      Jmat = zero
      ! (mu,nu^T) pairs
      mu_vec = (/zero,zero,zero/)
print *, "SIZE(TQvecs)", SIZE(TQvecs)
      DO t = 1, SIZE(TQvecs)
!         print *, t, SIZE(TQvecs)
         Tvec(:) = TQvecs(t)%Tvec(:) * lattice_constants(:)
         box(:) = (/1e10,1e10,1e10/)
         CALL pmd1_get_shell_pairs('I',basis,mu_vec,Tvec,box,mu_nu)

         ! (lambda^S,sigma^{S+Q}) pairs
         box(:) = (4.5d0)*lattice_constants(:)
         DO q = 1, SIZE(TQvecs)
            Qvec(:) = TQvecs(q)%Tvec(:) * lattice_constants(:)
            ! Loop over "infinity" of S vectors
            DO i = -iSmax, iSmax
               DO j = -iSmax, iSmax
                  DO k = -iSmax, iSmax
!IF (abs(i)<2 .AND. abs(j)<2 .AND. abs(k)<2) CYCLE
                     Svec(1) =  i * lattice_constants(1)
                     Svec(2) =  j * lattice_constants(2)
                     Svec(3) =  k * lattice_constants(3)
                     CALL pmd1_get_shell_pairs('I',basis,Svec,Svec+Qvec,box,ls)

                     ! Evaluate J_matrix contribution from this sh-quartet
!print *, "SIZE(mu_nu), SIZE(ls)", SIZE(mu_nu), SIZE(ls)
                     CALL pmd1_update_Jmat(dens(:,:,q),mu_nu,ls,Jmat(:,:,t))

                  END DO
               END DO
            END DO
         END DO
         energy = energy + dens(1,1,t) * Jmat(1,1,t)
!print *, dens(1,1,t) * Jmat(1,1,t)
      END DO
print *, "Energy = ", energy

      DEALLOCATE(mu_nu,ls)
      NULLIFY(mu_nu,ls)

   END SUBROUTINE pmd1_get_J_matrix

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_update_Jmat(dens,mu_nu,lam_sig,Jmat)

      USE md1_utils, ONLY: md1_get_prim_batch

      IMPLICIT NONE
      REAL(REALK),        INTENT(IN)  :: dens(:,:)
      TYPE(md1_sh_pairs), INTENT(IN)  :: mu_nu(:), lam_sig(:)
      REAL(REALK),        INTENT(OUT) :: Jmat(:,:)

      TYPE(md1_prim_batch), POINTER :: mn_batch(:), ls_batch(:)
      TYPE(md1_prim_batch), TARGET  :: tmp1(basis%Cmax2)
      TYPE(md1_prim_batch), TARGET  :: tmp2(basis%Cmax2)
      INTEGER(INTK) :: mn, ls, np
      REAL(REALK) :: g

      DO mn = 1, SIZE(mu_nu)
!         print *, "(mu,nu) shell: P%lo,hi", mu_nu%P%lo, mu_nu%P%hi
!         print *, "               Q%lo,hi", mu_nu%Q%lo, mu_nu%Q%hi
         ! Get (mu,nu) primitive data
         CALL md1_get_prim_batch(basis,mu_nu(mn),tmp1,np)
         IF (np > 1) CALL fmm_quit('contracted AOs NYI')
         mn_batch => tmp1(1:np)

         DO ls = 1, SIZE(lam_sig)
            ! Get (lambda,sigma) primitive data
            CALL md1_get_prim_batch(basis,lam_sig(ls),tmp2,np)
            IF (np > 1) CALL fmm_quit('contracted AOs NYI')
            ls_batch => tmp2(1:np)
      
!            print *, "mu_nu:"
!            print *, "        P   =", mn_batch(1)%P
!            print *, "        exp =", mn_batch(1)%exp
!            print *, "        Kab =", mn_batch(1)%KAB
!            print *, "        Cab =", mn_batch(1)%CAB
!            print *, "lam_sig:"
!            print *, "        P   =", ls_batch(1)%P
!            print *, "        exp =", ls_batch(1)%exp
!            print *, "        Kab =", ls_batch(1)%KAB
!            print *, "        Cab =", ls_batch(1)%CAB

            g = ssss_int(mn_batch(1),ls_batch(1))
            Jmat(1,1) = Jmat(1,1) + g * dens(1,1) 

         END DO
      END DO

      CONTAINS

         FUNCTION ssss_int(mn,ls)

         IMPLICIT NONE  
         TYPE(md1_prim_batch), INTENT(IN) :: mn, ls
         REAL(REALK) :: ssss_int
 
         REAL(REALK), PARAMETER :: PI = 3.14159265358979d0
         REAL(REALK) :: alpha, p,q, Kpq, U, Rpq

         p = mn%exp
         q = ls%exp
         alpha = one / (one/p + one/q )

         Kpq = mn%KAB * mn%CAB * ls%KAB * ls%CAB
         U = (PI/(SQRT(p*q)))**3 * Kpq
         Rpq = SQRT(DOT_PRODUCT((mn%P-ls%P),(mn%P-ls%P)))

         IF (Rpq > zero) THEN
             ssss_int = U * ERF(SQRT(alpha) * Rpq) / Rpq
         ELSE
             ssss_int = two * U * SQRT(alpha/PI)
         END IF
!print *, Rpq, ssss_int, U/Rpq - ssss_int


         END FUNCTION ssss_int

   END SUBROUTINE pmd1_update_Jmat

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_write_multipole_integrals(LMAX,TQvecs)

      USE md1_multipole_ints, ONLY: md1_init_multipole_ints,            &
                                    md1_free_multipole_ints

      IMPLICIT NONE  
      INTEGER(INTK),      INTENT(IN) :: LMAX
      TYPE(T_vector_map), INTENT(IN) :: TQvecs(:)
      INTEGER(INTK) :: nTmoms, nSQmoms

      ! Build and store GTO multipole moments on disk
      WRITE(IO1,*), "PMD1 multipole integral LMAX =", LMAX

      ! Open multipole file for interface with pfmmlib
      OPEN(UNIT=IO2, FILE='multipoles.md1', STATUS='REPLACE',  &
           ACCESS='SEQUENTIAL', FORM='UNFORMATTED')
      REWIND(IO2)
                                    
      ! Write overlap matrix to disk for diagnostics 
      OPEN(UNIT=IO3, FILE="overlap_matrix.md1", STATUS='REPLACE',  &
           ACCESS='SEQUENTIAL', FORM='FORMATTED')
      REWIND(IO3)

      CALL md1_init_multipole_ints(basis,LMAX)
      CALL pmd1_build_and_write_periodic_mpoles(LMAX,TQvecs,nTmoms,nSQmoms)
      CALL md1_free_multipole_ints

      ! Mark end of file with negative L-index
      WRITE(IO2) 0,-1,0,0,0, 0d0,0d0,0d0, 0d0
      CLOSE(UNIT=IO2, STATUS='KEEP')
      WRITE(IO3,'(3I8,F25.15)') 0,-1,0, 0d0
      CLOSE(UNIT=IO3, STATUS='KEEP')

      ! Write header file with basic constants and dimensions
      OPEN(UNIT=IO2, FILE='multipoles_header.md1', STATUS='REPLACE',   &
           ACCESS='SEQUENTIAL', FORM='UNFORMATTED')
      WRITE(IO2) LMAX, basis%nbas, nTmoms, nSQmoms
      CLOSE(UNIT=IO2, STATUS='KEEP')

   END SUBROUTINE pmd1_write_multipole_integrals

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_build_and_write_periodic_mpoles(LMAX,TQvecs,nTmoms,nSQmoms)

      IMPLICIT NONE  
      INTEGER(INTK),      INTENT(IN)  :: LMAX
      TYPE(T_vector_map), INTENT(IN)  :: TQvecs(:)
      INTEGER(INTK),      INTENT(OUT) :: nTmoms, nSQmoms

      CALL pmd1_build_and_write_mu_nu(LMAX,TQvecs,nTmoms)
      CALL pmd1_build_and_write_lambda_sigma(LMAX,TQvecs,nSQmoms)

   END SUBROUTINE pmd1_build_and_write_periodic_mpoles

!-------------------------------------------------------------------------------

   SUBROUTINE pmd1_build_and_write_mu_nu(LMAX,TQvecs,nTmoms)

      USE pmd1_shell_pairs,   ONLY: pmd1_get_shell_pairs
      USE md1_multipole_ints, ONLY: md1_build_and_write_sph_mpoles

      IMPLICIT NONE  
      INTEGER(INTK),      INTENT(IN)  :: LMAX
      TYPE(T_vector_map), INTENT(IN)  :: TQvecs(:)
      INTEGER(INTK),      INTENT(OUT) :: nTmoms

      TYPE(md1_sh_pairs), POINTER :: sh_pairs(:)
      INTEGER(INTK) :: i,j, nmoms
      REAL(REALK) :: mu_vec(3), nu_vec(3), box(3)

      ! Store (mu,nu) pairs with lattice shift T
      WRITE(IO1,*) "Building (mu,nu^T) moments."
      WRITE(IO2) 0,-2,0,0,0, 0d0,0d0,0d0, 0d0

      nTmoms = 0
      mu_vec = (/zero,zero,zero/)
      ! Do not impose restriction that (mu,nu^T) must be centered in 0 cell
      box(:) = (/1e10,1e10,1e10/)

      DO i = 1, SIZE(TQvecs)

         nu_vec(:) = TQvecs(i)%Tvec(:) * lattice_constants(:)
         WRITE(IO2) i, TQvecs(i)%Tvec(:)
         ! Get non-zero shell pairs for this (mu,nu^T)
         CALL pmd1_get_shell_pairs('I',basis,mu_vec,nu_vec,box,sh_pairs)
         CALL md1_build_and_write_sph_mpoles(basis,LMAX,sh_pairs,nmoms)
         nTmoms = nTmoms + nmoms

         IF (i < SIZE(TQvecs)) THEN
            ! Write key to tell PFMM parser that another batch is coming
            WRITE(IO2) 0,-2,0,0,0, 0d0,0d0,0d0, 0d0
         ELSE
            ! Mark end of (mu,nu) pairs
            WRITE(IO2) 0,-22,0,0,0, 0d0,0d0,0d0, 0d0
         END IF

      END DO

      DEALLOCATE(sh_pairs)
      NULLIFY(sh_pairs)

   END SUBROUTINE pmd1_build_and_write_mu_nu

!-------------------------------------------------------------------------------
! Note we include (S,S+Q) pairs "both ways".
! This should not give weird doublecoutning as long as periodic
! density matrix reflects it.  (Same issue in non-periodic case with 
! (1,2) and (2,1) AO-pairs being counted as two or one unit.)

!FIXME!!!!!!!!!!!!
! THIS ALGORITHM DOESN'T WORK IF AO-PAIRS FALL EXACTLY ON THE BOX BOUNDARIES!!!
! FOR TESTING WE AVOID THIS ISSUE BY PUTTING THE AO'S OFF-CENTRE in basis.pfmm

   SUBROUTINE pmd1_build_and_write_lambda_sigma(LMAX,TQvecs,nSQmoms)

      USE pmd1_shell_pairs,   ONLY: pmd1_get_shell_pairs
      USE md1_multipole_ints, ONLY: md1_build_and_write_sph_mpoles

      IMPLICIT NONE  
      INTEGER(INTK),      INTENT(IN)  :: LMAX
      TYPE(T_vector_map), INTENT(IN)  :: TQvecs(:)
      INTEGER(INTK),      INTENT(OUT) :: nSQmoms

      TYPE(md1_sh_pairs), POINTER :: sh_pairs(:)
      INTEGER(INTK) :: i,j, nmoms
      REAL(REALK) :: Svec(3), Qvec(3), box(3)

      ! Store (lambda,sigma) pairs (in one box) with shift Q
      WRITE(IO1,*) "Building (lambda^S,sigma^(S+Q)) moments."
      WRITE(IO2) 0,-3,0,0,0, 0d0,0d0,0d0, 0d0

      ! Assume reference unit cell is centered at (0,0,0)
      box(:) = half*lattice_constants(:)

      nSQmoms = 0
      ! Loop over Q-vectors
      DO i = 1, SIZE(TQvecs)
         Qvec(:) = TQvecs(i)%Tvec(:) * lattice_constants(:)
         WRITE(IO2) i, TQvecs(i)%Tvec(:)
!print '(A,3F15.5)', "Qvec        = ", Qvec(:)

         ! Loop over S-vectors
         ! We take a conservative sample of S vectors = TQvecs
         ! (This should generate all non-zero pairs)
         nmoms = 0
         DO j = 1, SIZE(TQvecs)
            Svec(:) = TQvecs(j)%Tvec(:) * lattice_constants(:)
!print '(A,3F15.5)', "Svec+Qvec   = ", Svec(:)+Qvec(:)
            CALL pmd1_get_shell_pairs('I',basis,Svec,Svec+Qvec,box,sh_pairs)
            CALL md1_build_and_write_sph_mpoles(basis,LMAX,sh_pairs,nmoms)
            nSQmoms = nSQmoms + nmoms
!print *, i, j, nmoms
         END DO

         IF (i < SIZE(TQvecs)) THEN
            ! Write key to tell PFMM parser that another SQ-batch is coming
            WRITE(IO2) 0,-3,0,0,0, 0d0,0d0,0d0, 0d0
         ELSE
            ! Mark end of all (lambda,sigma) pairs
            WRITE(IO2) 0,-33,0,0,0, 0d0,0d0,0d0, 0d0
         END IF
      END DO

      DEALLOCATE(sh_pairs)
      NULLIFY(sh_pairs)

   END SUBROUTINE pmd1_build_and_write_lambda_sigma

!-------------------------------------------------------------------------------

END MODULE pmd1_main
